# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/24 0:09
# @Author  : Dell
# @File    : rag_base_generation.py
# @Software: PyCharm
# @Desc    : RAG 生成：模版化生成

from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain.prompts import ChatPromptTemplate
# 查询提示词模板库
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
llm = ChatOpenAI(model_name='qwen-plus', base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", temperature=0.5,
                 api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37")
# --------------简单的一段话-----------------------
# 1、加载文本
loader = TextLoader(file_path="xiyouji.txt",encoding="utf-8")
blog_docs = loader.load()
# 2、分割文档
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=300,chunk_overlap=50)
splits = text_splitter.split_documents(blog_docs)
# 3、构建一个支持语义搜索的向量数据库（数据库在内存中）
embd=DashScopeEmbeddings(
    model="text-embedding-v1", dashscope_api_key="sk-7b4f290b1a3e42a1a9a1957fa44eff37"
)
vectorstore = Chroma.from_documents(documents=splits,embedding=embd)
#4、检索
# 将一个向量数据库（vectorstore）转换为一个检索器（retriever），并设置检索参数：返回最相似的5个结果。
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
docs = retriever.get_relevant_documents("孙悟空的出生地在哪里？")
# 5、生成
template = """您是问答任务的助手。使用以下检索到的上下文来回答问题。如果您不知道答案，就说您不知道。
{context}

问题：{question}
"""
prompt = ChatPromptTemplate.from_template(template)
# --------------方式一：采用简单链-------------------------------

chain = prompt | llm
re=chain.invoke({"context":retriever,"question":"孙悟空的出生地在哪里？"})
print(re.content)

# -------------------方式二采用rag链--------------------
"""
prompt_hub_rag = hub.pull("rlm/rag-prompt") # 从Hub拉取模板
prompt = ChatPromptTemplate.from_template(template)
# 用户问题 → { 同时进行: context: 检索器获取相关文档,question: 直接传递问题} → prompt模板 → llm → 输出解析器 → 最终答案
rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)
re=rag_chain.invoke("孙悟空的出生地在哪里？")
print(re)
"""